from flask import Flask, request, jsonify
import subprocess
from ariadne import ObjectType, QueryType, MutationType, gql, make_executable_schema
from ariadne.asgi import GraphQL
from flask import render_template_string

app = Flask(__name__)

PLAYGROUND_HTML = '''
<!DOCTYPE html>
<html>

<head>
  <meta charset=utf-8/>
  <meta name="viewport" content="user-scalable=no, initial-scale=1.0, minimum-scale=1.0, maximum-scale=1.0, minimal-ui">
  <title>GraphQL Playground</title>
  <link rel="stylesheet" href="//cdn.jsdelivr.net/npm/graphql-playground-react/build/static/css/index.css" />
  <link rel="shortcut icon" href="//cdn.jsdelivr.net/npm/graphql-playground-react/build/favicon.png" />
  <script src="//cdn.jsdelivr.net/npm/graphql-playground-react/build/static/js/middleware.js"></script>
</head>

<body>
  <div id="root">
    <style>
      body {
        background-color: rgb(23, 42, 58);
        font-family: Open Sans, sans-serif;
        height: 90vh;
      }

      #root {
        height: 100%;
        width: 100%;
        display: flex;
        align-items: center;
        justify-content: center;
      }

      .loading {
        font-size: 32px;
        font-weight: 200;
        color: rgba(255, 255, 255, .6);
        margin-left: 20px;
      }

      img {
        width: 78px;
        height: 78px;
      }

      .title {
        font-weight: 400;
      }
    </style>
    <img src='//cdn.jsdelivr.net/npm/graphql-playground-react/build/logo.png' alt=''>
    <div class="loading"> Loading
      <span class="title">GraphQL Playground</span>
    </div>
  </div>
  <script>window.addEventListener('load', function (event) {
      GraphQLPlayground.init(document.getElementById('root'), {
        endpoint: '/graphql'
      })
    })</script>
</body>

</html>
'''

# REST API routes
@app.route('/rest/generate', methods=['POST'])
def rest_generate():
    return generate_text(request.json['prompt'])

@app.route('/rest/perplexity', methods=['POST'])
def rest_perplexity():
    return calculate_perplexity(request.json['text'])

@app.route('/rest/retrieval', methods=['POST'])
def rest_retrieval():
    return perform_retrieval(request.json['query'])

# Helper functions
def generate_text(prompt, number_cha=512):
    command = [
        "../llama-cli",
        "-m", "/mnt/hgfs/cloudapp/qwen2-1_5b-instruct-q3_k_m.gguf",
        "-p", prompt,
        "-n", str(number_cha),
    ]
    result = subprocess.run(command, capture_output=True, text=True)
    return jsonify({"response": result.stdout})

def calculate_perplexity(text):
    command = [
        "../llama-cli",
        "-m", "/mnt/hgfs/cloudapp/qwen2-1_5b-instruct-q3_k_m.gguf",
        "--perplexity",
        "-p", text,
    ]
    result = subprocess.run(command, capture_output=True, text=True)
    # Parse the perplexity value from the output
    perplexity = float(result.stdout.strip().split()[-1])
    return jsonify({"perplexity": perplexity})

def perform_retrieval(query):
    command = [
        "../llama-cli",
        "-m", "/mnt/hgfs/cloudapp/qwen2-1_5b-instruct-q3_k_m.gguf",
        "--context-file", "your_context_file.txt",
        "-p", query,
    ]
    result = subprocess.run(command, capture_output=True, text=True)
    return jsonify({"retrieval_result": result.stdout})

# GraphQL schema and resolvers
type_defs = gql("""
    type Query {
        generate(prompt: String!): GenerationResult!
        perplexity(text: String!): PerplexityResult!
        retrieval(query: String!): RetrievalResult!
    }

    type GenerationResult {
        response: String!
    }

    type PerplexityResult {
        perplexity: Float!
    }

    type RetrievalResult {
        retrievalResult: String!
    }
""")

query = QueryType()

@query.field("generate")
def resolve_generate(_, info, prompt):
    result = generate_text(prompt)
    return {"response": result.json["response"]}

@query.field("perplexity")
def resolve_perplexity(_, info, text):
    result = calculate_perplexity(text)
    return {"perplexity": result.json["perplexity"]}

@query.field("retrieval")
def resolve_retrieval(_, info, query):
    result = perform_retrieval(query)
    return {"retrievalResult": result.json["retrieval_result"]}

schema = make_executable_schema(type_defs, query)

@app.route("/graphql", methods=["GET"])
def graphql_playground():
    return PLAYGROUND_HTML, 200

@app.route("/graphql", methods=["POST"])
def graphql_server():
    data = request.get_json()
    success, result = graphql_sync(
        schema,
        data,
        context_value=request,
        debug=app.debug
    )
    status_code = 200 if success else 400
    return jsonify(result), status_code

# Simple HTML form for testing
@app.route('/')
def index():
    html = """
    <!DOCTYPE html>
    <html>
    <head>
        <title>LLM API Test</title>
    </head>
    <body>
        <h1>REST API Test</h1>
        <h2>Generate Text</h2>
        <form id="restGenerateForm">
            <textarea name="prompt" rows="4" cols="50"></textarea><br>
            <button type="submit">Generate (REST)</button>
        </form>
        <div id="restGenerateResult"></div>

        <h2>Calculate Perplexity</h2>
        <form id="restPerplexityForm">
            <textarea name="text" rows="4" cols="50"></textarea><br>
            <button type="submit">Calculate Perplexity (REST)</button>
        </form>
        <div id="restPerplexityResult"></div>

        <h2>Perform Retrieval</h2>
        <form id="restRetrievalForm">
            <input type="text" name="query"><br>
            <button type="submit">Retrieve (REST)</button>
        </form>
        <div id="restRetrievalResult"></div>

        <h1>GraphQL API Test</h1>
        <h2>Generate Text</h2>
        <form id="graphqlGenerateForm">
            <textarea name="prompt" rows="4" cols="50"></textarea><br>
            <button type="submit">Generate (GraphQL)</button>
        </form>
        <div id="graphqlGenerateResult"></div>

        <h2>Calculate Perplexity</h2>
        <form id="graphqlPerplexityForm">
            <textarea name="text" rows="4" cols="50"></textarea><br>
            <button type="submit">Calculate Perplexity (GraphQL)</button>
        </form>
        <div id="graphqlPerplexityResult"></div>

        <h2>Perform Retrieval</h2>
        <form id="graphqlRetrievalForm">
            <input type="text" name="query"><br>
            <button type="submit">Retrieve (GraphQL)</button>
        </form>
        <div id="graphqlRetrievalResult"></div>

        <script>
            // REST API calls
            document.getElementById('restGenerateForm').onsubmit = function(e) {
                e.preventDefault();
                fetch('/rest/generate', {
                    method: 'POST',
                    headers: {'Content-Type': 'application/json'},
                    body: JSON.stringify({prompt: this.prompt.value})
                })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('restGenerateResult').textContent = data.response;
                });
            };

            document.getElementById('restPerplexityForm').onsubmit = function(e) {
                e.preventDefault();
                fetch('/rest/perplexity', {
                    method: 'POST',
                    headers: {'Content-Type': 'application/json'},
                    body: JSON.stringify({text: this.text.value})
                })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('restPerplexityResult').textContent = 'Perplexity: ' + data.perplexity;
                });
            };

            document.getElementById('restRetrievalForm').onsubmit = function(e) {
                e.preventDefault();
                fetch('/rest/retrieval', {
                    method: 'POST',
                    headers: {'Content-Type': 'application/json'},
                    body: JSON.stringify({query: this.query.value})
                })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('restRetrievalResult').textContent = data.retrieval_result;
                });
            };

            // GraphQL API calls
            document.getElementById('graphqlGenerateForm').onsubmit = function(e) {
                e.preventDefault();
                fetch('/graphql', {
                    method: 'POST',
                    headers: {'Content-Type': 'application/json'},
                    body: JSON.stringify({
                        query: `query($prompt: String!) {
                            generate(prompt: $prompt) {
                                response
                            }
                        }`,
                        variables: {prompt: this.prompt.value}
                    })
                })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('graphqlGenerateResult').textContent = data.data.generate.response;
                });
            };

            document.getElementById('graphqlPerplexityForm').onsubmit = function(e) {
                e.preventDefault();
                fetch('/graphql', {
                    method: 'POST',
                    headers: {'Content-Type': 'application/json'},
                    body: JSON.stringify({
                        query: `query($text: String!) {
                            perplexity(text: $text) {
                                perplexity
                            }
                        }`,
                        variables: {text: this.text.value}
                    })
                })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('graphqlPerplexityResult').textContent = 'Perplexity: ' + data.data.perplexity.perplexity;
                });
            };

            document.getElementById('graphqlRetrievalForm').onsubmit = function(e) {
                e.preventDefault();
                fetch('/graphql', {
                    method: 'POST',
                    headers: {'Content-Type': 'application/json'},
                    body: JSON.stringify({
                        query: `query($query: String!) {
                            retrieval(query: $query) {
                                retrievalResult
                            }
                        }`,
                        variables: {query: this.query.value}
                    })
                })
                .then(response => response.json())
                .then(data => {
                    document.getElementById('graphqlRetrievalResult').textContent = data.data.retrieval.retrievalResult;
                });
            };
        </script>
    </body>
    </html>
    """
    return render_template_string(html)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5001, debug=True)